#include "limbo/ip/state.hpp"
#include "limbo/stack.hpp"
#include "limbo/udp/receiver.hpp"
#include "test-utils.h"
#include <type_traits>

using namespace limbo;

using IP = ip::State<void, ip::Opts::calculate_checksum>;
using UDP = udp::State<IP>;
using Receiver = udp::Receiver<UDP>;
using MyStack = Stack<Layer<Receiver>, Layer<UDP>, Layer<IP>>;

unsigned char example_raw[20 + 8 + 5] = {
    /* hdr */ 0x45, 0x00, 0x00, 0x21, 0x00, 0x05, 0x00, 0x00,
    /* ttl */ 0x06, 0x11, 0xaa, 0x50, 0x0a, 0x00, 0x1f, 0x7c,
    /* dst */ 0xe0, 0x00, 0x00, 0xfb, 0x00, 0x07, 0x00, 0x08,
    /* len */ 0x00, 0x0D, 0x00, 0x00, 'h',  'e',  'l',  'l',  'o'};
auto example_chunk = Chunk(example_raw, sizeof(example_raw));
auto payload_chunk = Chunk((void *)(example_raw + 28), 5);

auto src = make_address("10.0.31.124");
auto dst = make_address("224.0.0.251");

TEST_CASE("send", "[udp-receiver]") {
  auto stack = MyStack();
  auto &ip_state = stack.get<2, 0>();
  auto &udp_recv = stack.get<0, 0>();
  ip_state.init(0x04);
  udp_recv.init({src, 7}, 6);

  unsigned char buff[sizeof(example_raw)];
  auto buff_chunk = Chunk(buff, sizeof(buff));
  auto result = udp_recv.send(stack, buff_chunk, {dst, 8}, payload_chunk);

  REQUIRE(result);
  auto res_buff = result.consumed();
  CHECK((void *)res_buff.data() == (void *)buff);
  CHECK(res_buff.size() == example_chunk.size());
  CHECK(res_buff == example_chunk);
}

TEST_CASE("receive", "[udp-receiver]") {
  auto stack = MyStack();
  auto &udp_recv = stack.get<0, 0>();
  udp_recv.init({dst, 8}, 6); /* dst and src are swapped */
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &udp_recv);
  auto &p = udp_recv.get_parsed();
  CHECK(p.source_port == 7);
  CHECK(p.dest_port == 8);
  CHECK(p.length == payload_chunk.size() + 8);
  CHECK(p.payload == payload_chunk);
  CHECK(p.container->source == src);
  CHECK(p.container->destination == dst);
}

TEST_CASE("receive to wrong port", "[udp-receiver]") {
  auto stack = MyStack();
  auto &udp_recv = stack.get<0, 0>();
  auto &udp_state = stack.get<1, 0>();
  /* dst and src are NOT swapped  =>
   * listening port 7, but receving upon port 8 */
  udp_recv.init({src, 7}, 6);
  auto result = stack.recv(example_chunk, nullptr);
  REQUIRE(result);
  REQUIRE(result.consumed() == example_chunk);
  REQUIRE(result.state() == &udp_state);
  auto &p = udp_state.get_parsed();
  CHECK(p.source_port == 7);
  CHECK(p.dest_port == 8);
  CHECK(p.length == payload_chunk.size() + 8);
  CHECK(p.payload == payload_chunk);
  CHECK(p.container->source == src);
  CHECK(p.container->destination == dst);
}
